import os
import shutil
import random
import cv2
import numpy as np
import albumentations as A
from PIL import Image
import uuid
from concurrent.futures import ProcessPoolExecutor
import functools
import logging

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

def count_photos_in_subfolders(parent_folder):
    subfolders = [f.path for f in os.scandir(parent_folder) if f.is_dir()]
    photo_counts = {subfolder: len([f for f in os.listdir(subfolder) if f.lower().endswith(('.png', '.jpg', '.jpeg'))]) for subfolder in subfolders}
    return photo_counts

def is_valid_image(file_path):
    try:
        with Image.open(file_path) as img:
            img.verify()
        return True
    except:
        return False

def enhance_photo(photo_path, new_subfolder_path):
    try:
        image = cv2.imread(photo_path)
        if image is None:
            logger.error(f"Unable to read image: {photo_path}")
            return None
        
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        height, width = image.shape[:2]

        transform = A.Compose([
            A.RandomRotate90(p=0.5),
            A.Flip(p=0.5),
            A.Transpose(p=0.5),
            A.OneOf([
                A.GaussNoise(var_limit=(10.0, 50.0)),
                A.MultiplicativeNoise(),
            ], p=0.2),
            A.OneOf([
                A.MotionBlur(p=0.2),
                A.MedianBlur(blur_limit=3, p=0.1),
                A.Blur(blur_limit=3, p=0.1),
                A.GaussianBlur(blur_limit=(3, 7), p=0.1),
            ], p=0.2),
            A.ShiftScaleRotate(shift_limit=0.0625, scale_limit=0.1, rotate_limit=45, p=0.2, border_mode=cv2.BORDER_CONSTANT),
            A.OneOf([
                A.OpticalDistortion(p=0.3, border_mode=cv2.BORDER_CONSTANT),
                A.GridDistortion(p=0.1, border_mode=cv2.BORDER_CONSTANT),
                A.ElasticTransform(p=0.2, border_mode=cv2.BORDER_CONSTANT),
            ], p=0.2),
            A.CoarseDropout(max_holes=8, max_height=max(height//16, 1), max_width=max(width//16, 1), min_holes=1, min_height=1, min_width=1, fill_value=0, p=0.2),
        ])

        augmented = transform(image=image)
        augmented_image = augmented['image']

        new_filename = f'enhanced_{uuid.uuid4()}.jpg'
        new_path = os.path.join(new_subfolder_path, new_filename)
        cv2.imwrite(new_path, cv2.cvtColor(augmented_image, cv2.COLOR_RGB2BGR))
        return new_filename
    except Exception as e:
        logger.error(f"Error enhancing {photo_path}: {str(e)}", exc_info=True)
        return None

def process_subfolder(subfolder, new_folder, target_photo_count):
    photos = [f for f in os.listdir(subfolder) if f.lower().endswith(('.png', '.jpg', '.jpeg'))]
    new_subfolder_path = os.path.join(new_folder, os.path.basename(subfolder))
    os.makedirs(new_subfolder_path, exist_ok=True)

    copied_photos = []

    # 复制所有原始图片
    for photo in photos:
        src_path = os.path.join(subfolder, photo)
        if is_valid_image(src_path):
            dest_path = os.path.join(new_subfolder_path, photo)
            shutil.copyfile(src_path, dest_path)
            copied_photos.append(photo)

    # 如果原始图片数量不足，则进行增强
    while len(copied_photos) < target_photo_count:
        photo = random.choice(photos)
        src_path = os.path.join(subfolder, photo)
        new_filename = enhance_photo(src_path, new_subfolder_path)
        if new_filename:
            copied_photos.append(new_filename)

    return len(copied_photos)

def copy_photos_to_new_folder(parent_folder, new_folder, photo_counts):
    max_photo_count = max(photo_counts.values())
    print(f'Max photos in a subfolder: {max_photo_count}')

    with ProcessPoolExecutor() as executor:
        process_func = functools.partial(process_subfolder, new_folder=new_folder, target_photo_count=max_photo_count)
        results = list(executor.map(process_func, photo_counts.keys()))

    for subfolder, count in zip(photo_counts.keys(), results):
        print(f'Subfolder: {subfolder}, Final copied and enhanced photos: {count}')

if __name__ == "__main__":
    parent_folder = 'train_ori_20240804'
    new_folder = 'imagenet/train'

    photo_counts = count_photos_in_subfolders(parent_folder)
    copy_photos_to_new_folder(parent_folder, new_folder, photo_counts)